﻿using Renci.SshNet;
using Renci.SshNet.Common;
using System.Collections.Concurrent;

namespace Tool.Utils
{
    public class ImprovedJumpHostConnectionManager : IDisposable
    {
        private readonly ConcurrentDictionary<string, SshClient> _activeConnections =
            new ConcurrentDictionary<string, SshClient>();

        private readonly ConcurrentDictionary<string, SshClient> _jumpClients =
            new ConcurrentDictionary<string, SshClient>();

        private readonly ConcurrentDictionary<string, ForwardedPortLocal> _activeTunnels =
            new ConcurrentDictionary<string, ForwardedPortLocal>();

        public class ConnectionConfig
        {
            public string JumpHost { get; set; }
            public int JumpPort { get; set; } = 22;
            public string JumpUser { get; set; }
            public string JumpPassword { get; set; }
            public PrivateKeyFile JumpPrivateKey { get; set; }
            public int JumpTimeoutSeconds { get; set; } = 30;

            public string TargetHost { get; set; }
            public int TargetPort { get; set; } = 22;
            public string TargetUser { get; set; }
            public string TargetPassword { get; set; }
            public PrivateKeyFile TargetPrivateKey { get; set; }
            public int TargetTimeoutSeconds { get; set; } = 30;
        }

        public async Task<SshClient> GetTargetConnectionAsync(ConnectionConfig config)
        {
            var connectionKey = $"{config.JumpHost}:{config.JumpPort}=>{config.TargetHost}:{config.TargetPort}";

            if (_activeConnections.TryGetValue(connectionKey, out var existingClient) && existingClient.IsConnected)
            {
                return existingClient;
            }

            try
            {
                // 建立跳板机连接
                var jumpClient = await CreateJumpClientAsync(config);
                _jumpClients[connectionKey] = jumpClient;

                // 创建隧道
                var localPort = new Random().Next(10000, 20000);
                var tunnel = new ForwardedPortLocal("127.0.0.1", (uint)localPort, config.TargetHost, (uint)config.TargetPort);
                jumpClient.AddForwardedPort(tunnel);

                // 启动隧道前检查跳板机连接状态
                if (!jumpClient.IsConnected)
                {
                    throw new InvalidOperationException("跳板机连接未建立");
                }

                tunnel.Start();

                // 通过隧道连接到目标服务器
                var targetClient = await CreateTargetClientAsync(config, localPort);

                // 存储连接信息
                _activeConnections[connectionKey] = targetClient;
                _activeTunnels[connectionKey] = tunnel;

                return targetClient;
            }
            catch (Exception)
            {
                // 清理已创建的资源
                if (_jumpClients.TryRemove(connectionKey, out var jumpClient))
                {
                    jumpClient.Disconnect();
                    jumpClient.Dispose();
                }
                throw;
            }
        }

        private async Task<SshClient> CreateJumpClientAsync(ConnectionConfig config)
        {
            return await Task.Run(() =>
            {
                AuthenticationMethod jumpAuth;
                if (config.JumpPrivateKey != null)
                {
                    jumpAuth = new PrivateKeyAuthenticationMethod(config.JumpUser, config.JumpPrivateKey);
                }
                else
                {
                    jumpAuth = new PasswordAuthenticationMethod(config.JumpUser, config.JumpPassword);
                }

                var jumpConnectionInfo = new Renci.SshNet.ConnectionInfo(config.JumpHost, config.JumpPort, config.JumpUser, jumpAuth)
                {
                    Timeout = TimeSpan.FromSeconds(config.JumpTimeoutSeconds)
                };

                var jumpClient = new SshClient(jumpConnectionInfo);
                jumpClient.Connect();

                // 验证连接是否成功
                if (!jumpClient.IsConnected)
                {
                    throw new SshConnectionException("无法连接到跳板机");
                }

                return jumpClient;
            });
        }

        private async Task<SshClient> CreateTargetClientAsync(ConnectionConfig config, int localPort)
        {
            return await Task.Run(() =>
            {
                AuthenticationMethod targetAuth;
                if (config.TargetPrivateKey != null)
                {
                    targetAuth = new PrivateKeyAuthenticationMethod(config.TargetUser, config.TargetPrivateKey);
                }
                else
                {
                    targetAuth = new PasswordAuthenticationMethod(config.TargetUser, config.TargetPassword);
                }

                var targetConnectionInfo = new Renci.SshNet.ConnectionInfo("127.0.0.1", localPort, config.TargetUser, targetAuth)
                {
                    Timeout = TimeSpan.FromSeconds(config.TargetTimeoutSeconds)
                };

                var targetClient = new SshClient(targetConnectionInfo);

                // 重试机制
                var maxRetries = 3;
                for (int i = 0; i < maxRetries; i++)
                {
                    try
                    {
                        targetClient.Connect();
                        if (targetClient.IsConnected)
                        {
                            return targetClient;
                        }
                    }
                    catch (SshConnectionException) when (i < maxRetries - 1)
                    {
                        Task.Delay(1000).Wait(); // 等待1秒后重试
                        continue;
                    }
                }

                throw new SshConnectionException("无法通过隧道连接到目标服务器");
            });
        }

        public void ReleaseConnection(string connectionKey)
        {
            if (_activeConnections.TryRemove(connectionKey, out var client))
            {
                try
                {
                    client.Disconnect();
                    client.Dispose();
                }
                catch { /* 忽略关闭连接时的异常 */ }
            }

            if (_activeTunnels.TryRemove(connectionKey, out var tunnel))
            {
                try
                {
                    tunnel.Stop();
                }
                catch { /* 忽略停止隧道时的异常 */ }
            }

            if (_jumpClients.TryRemove(connectionKey, out var jumpClient))
            {
                try
                {
                    jumpClient.Disconnect();
                    jumpClient.Dispose();
                }
                catch { /* 忽略关闭跳板机连接时的异常 */ }
            }
        }

        public void Dispose()
        {
            foreach (var kvp in _activeConnections)
            {
                try
                {
                    kvp.Value?.Disconnect();
                    kvp.Value?.Dispose();
                }
                catch { /* 忽略异常 */ }
            }
            _activeConnections.Clear();

            foreach (var kvp in _activeTunnels)
            {
                try
                {
                    kvp.Value?.Stop();
                }
                catch { /* 忽略异常 */ }
            }
            _activeTunnels.Clear();

            foreach (var kvp in _jumpClients)
            {
                try
                {
                    kvp.Value?.Disconnect();
                    kvp.Value?.Dispose();
                }
                catch { /* 忽略异常 */ }
            }
            _jumpClients.Clear();
        }
    }
}
